##########################################################################################

library(data.table)
library(optparse)
library(Seurat)
library(ArchR)
library(ggsci)
library(dplyr)
library(ComplexHeatmap)

##########################################################################################
option_list <- list(
    make_option(c("--comine_data_file"), type = "character"),
    make_option(c("--scriptPath"), type = "character"),
    make_option(c("--divide"), type = "character"),
    make_option(c("--out_path"), type = "character") 
)

if(1!=1){
    
    ## 整合atac和rna的文件
    comine_data_file <- "~/20231121_singleMuti/results/qc_atac_v2/all/testis_combined_peak.combineRNA.qc.Rdata"

    ## 既往研究整理的代码
    scriptPath <- "~/20231121_singleMuti/scripts/scScalpChromatin"

    divide <- 3

    ## 输出
    out_path <- "~/20231121_singleMuti/results/celltype_plot"

}


###########################################################################################
parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

comine_data_file <- opt$comine_data_file
scriptPath <- opt$scriptPath
out_path <- opt$out_path
divide <- as.numeric(opt$divide)

dir.create( out_path , recursive = T)
#dir.create( paste0(out_path , "/peakMA") , recursive = T)

##########################################################################################
## 已发表文献写好的脚本
source(paste0(scriptPath, "/plotting_config.R"))
source(paste0(scriptPath, "/misc_helpers.R"))
source(paste0(scriptPath, "/matrix_helpers.R"))
source(paste0(scriptPath, "/archr_helpers.R"))
source(paste0(scriptPath, "/GO_wrappers.R"))

##########################################################################################
 
a <- load(comine_data_file)
# testis_combined_peak_combineRNA

## 细胞顺序
cell_order <- c("SSC", "Differenting&Differented SPG", "Leptotene",
    "Zygotene", "Patchytene", "Diplotene",
    "Early stage of spermatids", "Round&ElongateS.tids", "Sperm",
    "Leydig cells", "Myoid cells", "Pericytes",
    "Sertoli cells", "Endothelial cells", "NKT cells", "Macrophages"
    )

## 细胞颜色
use_colors <- c(pal_npg("nrc")(10) , pal_jco("default")(6))
names(use_colors) <- c("Myoid cells" , "Leydig cells" , "Endothelial cells" , "Zygotene" , "Round&ElongateS.tids" , 
"Patchytene" , "SSC" , "Sperm" , "Diplotene" , "Early stage of spermatids" , "Leptotene" , 
"Sertoli cells" , "Macrophages" , "Differenting&Differented SPG" , "Pericytes" , "NKT cells" )

## 三类细胞
SSC_SPG <- c("SSC" , "Differenting&Differented SPG")
SPC <- c("Leptotene" , "Zygotene", "Patchytene",
    "Diplotene" , "Early stage of spermatids")
SPT <- c("Round&ElongateS.tids" , "Sperm")

##########################################################################################

if(divide == 3){
  ## 合并细胞为三大类
  testis_combined_peak_combineRNA@cellColData$cell_type <- ifelse(testis_combined_peak_combineRNA@cellColData$cell_type %in% SSC_SPG , "SSC_SPG" , testis_combined_peak_combineRNA@cellColData$cell_type)
  testis_combined_peak_combineRNA@cellColData$cell_type <- ifelse(testis_combined_peak_combineRNA@cellColData$cell_type %in% SPC , "SPC" , testis_combined_peak_combineRNA@cellColData$cell_type)
  testis_combined_peak_combineRNA@cellColData$cell_type <- ifelse(testis_combined_peak_combineRNA@cellColData$cell_type %in% SPT , "SPT" , testis_combined_peak_combineRNA@cellColData$cell_type)
  grp_order <- c("SSC_SPG" , "SPC" , "SPT")
  cell_order <- grp_order
  use_colors <- c(pal_npg("nrc")(3))
  names(use_colors) <- grp_order
  ## 不能变为factor，报错Error in cor(estbgdP, obsbgdP) : incompatible dimensions
  #testis_combined_peak_combineRNA@cellColData$cell_type <- factor( testis_combined_peak_combineRNA@cellColData$cell_type , levels = c("SSC_SPG" , "SPC" , "SPT") , order = T )
}else{
  #testis_combined_peak_combineRNA@cellColData$cell_type <- factor( testis_combined_peak_combineRNA@cellColData$cell_type , levels = grp_order[grp_order %in% unique(testis_combined_peak_combineRNA@cellColData$cell_type)] , order = T )
}

##########################################################################################
## 所有peak，包含peak的注释
atac_proj <- testis_combined_peak_combineRNA

all_peak <- data.frame(atac_proj@peakSet)
all_peak$cell_type <- sapply( strsplit(all_peak$GroupReplicate , "._.") , "[" , 1 )
all_peak$peakName <- paste0( all_peak$seqnames , ":" , all_peak$start , "-" , all_peak$end )

out_file <- paste0(out_path , "/AllPeak.tsv" )
write.table( all_peak , out_file , row.names = F , quote = F , sep = "\t" )

##########################################################################################
## Identify Marker Peaks while controling for TSS and Depth Biases
##########################################################################################
##  we tell ArchR to account for differences in data quality amongst the cell groups by setting the bias parameter 
## to account for TSS enrichment and the number of unique fragments per cell.

use_groups <- unique(atac_proj$cell_type)
cell_order <- cell_order[cell_order %in% use_groups]

markerPeaks <- getMarkerFeatures(
    ArchRProj = atac_proj, 
    useMatrix = "PeakMatrix", 
    groupBy = "cell_type",
    useGroups = use_groups,
    bias = c("TSSEnrichment", "log10(nFrags)"),
    testMethod = "wilcoxon"
)

## 输出peak的对比
markerList <- getMarkers(markerPeaks, cutOff = "FDR <= 0.1 & Log2FC >= 0.5")
out_file <- paste0(out_path , "/plotMarkerPeakHeatmap.tsv" )
markerList <- data.frame(markerList)
markerList$peakName <- paste0( markerList$seqnames , ":" , markerList$start , "-" , markerList$end )
markerList <- merge( markerList , unique(all_peak[,c("peakType", "peakName")]) , by = "peakName" )
write.table( markerList , out_file , row.names = F , quote = F , sep = "\t" )

#Visualize Markers as a heatmap
heatmapPeaks <- plotMarkerHeatmap(
  seMarker = markerPeaks[,cell_order], 
  cutOff = "FDR <= 0.1 & Log2FC >= 0.5",
  nLabel = 1, # It still seems like there's not actually a way to NOT plot any labels
  binaryClusterRows = TRUE,
  clusterCols = FALSE,
  transpose = FALSE
)

out_file <- paste0( out_path , "/plotMarkerPeakHeatmap.pdf" ) 
pdf(out_file , height = 10 , width = 7)
draw(heatmapPeaks, heatmap_legend_side="bot", annotation_legend_side="bot")
dev.off()

## Marker Peak MA and Volcano Plots
#for(cellType in colnames(markerPeaks)){
#    pma <- markerPlot(seMarker = markerPeaks, name = cellType, cutOff = "FDR <= 0.1 & Log2FC >= 0.5", plotAs = "MA")
#    out_file <- paste0( out_path , "/peakMA/" , cellType , "_markerPlot.pdf" ) 
#    pdf(out_file , height = 10 , width = 7)
#    print(pma)
#    dev.off()
#}


##########################################################################################
## 所有motif所在的peak
if(1!=1){
    motifPositions <- getMatches(atac_proj, name="Motif")

    all_motif_match <- c()
    for(motif in colnames(motifPositions)){

        print(motif)
        match_pos <- which(assay(motifPositions)[,motif])
        tmp_peak <- all_peak[match_pos,]
        tmp_peak$motif_in <- motif

        all_motif_match <- bind_rows( all_motif_match , tmp_peak )
    }

    out_file <- paste0( out_path , "/AllPeak.containMotif.tsv" ) 
    write.table( all_motif_match , out_file , row.names = F , quote = F , sep = "\t" )

    ## 提取所有motif的位置
    motifPositions <- getPositions(atac_proj, name="Motif")
    out_file <- paste0( out_path , "/AllMotif.tsv" ) 
    write.table( motifPositions , out_file , row.names = F , quote = F , sep = "\t" )
}

##########################################################################################
## Motif Enrichments
##########################################################################################
# Identify Motif Enrichments
enrichMotifs <- peakAnnoEnrichment(
    seMarker = markerPeaks,
    ArchRProj = atac_proj,
    peakAnnotation = "Motif",
    cutOff = "FDR <= 0.1 & Log2FC >= 0.5"
  )

# Rename motifs for more aesthetic plotting:
rownames(enrichMotifs) <- lapply(rownames(enrichMotifs), function(x) strsplit(x, "_")[[1]][1]) %>% unlist()
df <- data.frame(TF = rownames(enrichMotifs), mlog10Padj = assay(enrichMotifs)[,1])

# Subset to clusters that have at least some enrichment
log10pCut <- 10

#ArchR Heatmap
heatmapEM <- plotEnrichHeatmap(
    enrichMotifs[,cell_order], 
    transpose=FALSE
)

out_file <- paste0( out_path , "/plotMotifEnrichHeatmap.pdf" ) 
pdf(out_file , height = 15 , width = 7)
draw(heatmapEM, heatmap_legend_side="bot", annotation_legend_side="bot")
dev.off()

#### 下面是好看的
# Let's plot a different style heatmap
plot_mat <- plotEnrichHeatmap(enrichMotifs[,cell_order], transpose=FALSE, 
  cutOff=log10pCut, returnMatrix=TRUE)

plot_mat <- plot_mat[!grepl("^ENSG", rownames(plot_mat)),]
rownames(plot_mat) <- sapply( strsplit(rownames(plot_mat) , "[( ]" , ) , "[" , 1)

out_file <- paste0(out_path, "/Motifs-Enriched-Heatmap-CellType.tsv")
write.table( plot_mat , out_file , row.names = T , quote = F , sep = "\t" )

# Let's plot a different style heatmap
## 所有的motif
plot_mat_all <- plotEnrichHeatmap(enrichMotifs[,cell_order], transpose=FALSE, 
  cutOff=log10pCut, returnMatrix=TRUE , n = 10000)
plot_mat_all <- plot_mat_all[!grepl("^ENSG", rownames(plot_mat_all)),]
rownames(plot_mat_all) <- sapply( strsplit(rownames(plot_mat_all) , "[( ]" , ) , "[" , 1)
out_file <- paste0(out_path, "/Motifs-Enriched-Heatmap-CellType.all.tsv")
write.table( plot_mat_all , out_file , row.names = T , quote = F , sep = "\t" )

# Get colors for cluster annotation
# (Link FineClusts to BroadClust cmap)
cM <- as.matrix(confusionMatrix(atac_proj$cell_type, atac_proj$cell_type))
map_colors <- apply(cM, 1, function(x)colnames(cM)[which.max(x)])
plotColors <- map_colors[colnames(plot_mat)]

bc_to_nc_map <- invertList(plotColors)
plot_mat <- prettyOrderMat(plot_mat[,cell_order], clusterCols=FALSE, cutOff=1)$mat

fontsize <- 6
ht_opt$simple_anno_size <- unit(0.25, "cm")
ta <- HeatmapAnnotation(
    atac_cluster=plotColors[colnames(plot_mat)],
    col=list(atac_cluster=use_colors), 
    show_legend=c(atac_cluster=FALSE), 
    show_annotation_name = c(atac_cluster=FALSE))

pdf(paste0(out_path, "/Motifs-Enriched-Heatmap-CellType.pdf"), width=6, height=10)
hm <- BORHeatmap(
  plot_mat, 
  limits=c(0.0,100.0), 
  clusterCols=FALSE, clusterRows=FALSE,
  labelCols=TRUE, labelRows=TRUE,
  dataColors = cmaps_BOR$comet,
  bottom_annotation = ta,
  row_names_side = "left",
  row_names_gp = gpar(fontsize = fontsize),
  column_names_gp = gpar(fontsize = fontsize),
  width = ncol(plot_mat)*unit(0.4, "cm"),
  height = nrow(plot_mat)*unit(0.2, "cm"),
  legendTitle="Norm.Enrichment -log10(P-adj)[0-Max]",
  border_gp = gpar(col="black") # Add a black border to entire heatmap
  )
draw(hm)
dev.off()
